import random
from math import pi
from typing import List
from Modules.foward_kinematics import aubo_FK
import numpy.linalg as lg

from Modules.auboi5_collision_check import auboi5_collision_check
from Modules.robot_class import Cube_obstacle, Joint_point,Point

import time

class RRT:
    """
    * 快速扩展随机树算法(Rapidly Exploring Random Tree),
    * 它的想法就是从根结点长出一棵树当树枝长到终点的时候这样就能找到从终点到根节点的唯一路径
    * 1. 进行随机采样在空间中随机选择一个点Xrand,寻找距离Xrand最近的节点（如果是第一次那就是初始节点）
    * 2. 随后我们进行树的生长，将Xnear--Xrand连接起来作为树生长的方向，设置步长作为树枝的长度，就会产生Xnew节点
    * 3. 将树枝节点和根结点加在一起作为新节点。
    """
    def __init__(self, start: Joint_point, goal: Joint_point, cube_obstacle: Cube_obstacle,
                 joint_limit_angle: List[Joint_point], num_iterations: int,
                 plan_area=[Joint_point([-pi] * 6), Joint_point([pi] * 6)]) -> None:

        self.start = start  # 机器人关节角起点
        self.goal = goal    # 机器人关节角终点
        self.cube_obstacle = cube_obstacle  # 障碍物
        self.joint_limit_angle = joint_limit_angle  # 机器人关节限制角度
        # joint_limit_angle = [JointPoint([[0] * 6]), JointPoint([2*pi] * 6)]
        self.plan_area = plan_area  # 关节角空间，默认-pi到pi

        # Set the step length, 默认5°
        self.step_length = 10 * pi / 180

        # Set the possibility of set the goal as new joint point,概率0-1
        self.set_goal = 0.1

        # Store the joint_point
        self.joint_point_list = [self.start]

        # Set the number of iterations,迭代次数
        self.num_iterations = num_iterations

        # 桌子建模需要
        self.limit_position = float(0.05)

    def generate_new_joint_point(self) -> Joint_point:
        """
        * 随机生成6个关节角度值
        """
        theta1 = random.uniform(self.plan_area[0].theta[0], self.plan_area[1].theta[0])
        theta2 = random.uniform(self.plan_area[0].theta[1], self.plan_area[1].theta[1])
        theta3 = random.uniform(self.plan_area[0].theta[2], self.plan_area[1].theta[2])
        theta4 = random.uniform(self.plan_area[0].theta[3], self.plan_area[1].theta[3])
        theta5 = random.uniform(self.plan_area[0].theta[4], self.plan_area[1].theta[4])
        theta6 = 0
        return Joint_point([theta1, theta2, theta3, theta4, theta5, theta6])

    def get_the_nearest_joint_point(self, new_random_joint_point: Joint_point) -> int:
        """
        * Get the nearest joint point in self.joint_point_list
        * 创建一个空列表d，用于存储new_random_joint_point与self.joint_point_list中每个点的距离。
        * 然后，使用for循环遍历self.joint_point_list中的每个点，计算它们与new_random_joint_point之间的距离
        *（使用lg.norm()函数计算向量的范数），并将结果添加到列表d中。最后，使用min()函数找到列表d中的最小值，
        * 并使用index()方法获取该最小值在列表中的索引，将其作为函数的返回值。
        """
        d = list()
        for joint_point in self.joint_point_list:
            d.append(lg.norm(joint_point.vector - new_random_joint_point.vector))
        return d.index(min(d))

    def expand_new_joint_point(self, nearest_joint_point: Joint_point,
                               new_random_joint_point: Joint_point) -> Joint_point:
        new_point = self.step_length * ((new_random_joint_point.vector - nearest_joint_point.vector) + (
                self.goal.vector - nearest_joint_point.vector)) + nearest_joint_point.vector
        return Joint_point(new_point.tolist())

    def collision_check(self, new_joint_point: Joint_point) -> bool:
        return auboi5_collision_check(new_joint_point, self.cube_obstacle)

    def get_farest_joint_point_from_start(self) -> Joint_point:
        d = list()
        for joint_point in self.joint_point_list:
            d.append(lg.norm(joint_point.vector - self.start.vector))
        return self.joint_point_list[d.index(max(d))]

    def check_joint_point_exceed_limit(self, joint_point: Joint_point) -> bool:
        """
        * Check the joint_point whether exceed the joint limit angle
        * 检查关节点是否超过了关节限制角度,是否小于最小关节角self.joint_limit_angle[0]，是否超出最大关节角self.joint_limit_angle[1]
        """
        exceed_limit = False
        for i in range(len(joint_point.theta)):
            if joint_point.theta[i] < self.joint_limit_angle[0].theta[i] or joint_point.theta[i] > \
                    self.joint_limit_angle[1].theta[i]:
                exceed_limit = True
                break

        # 步骤2：对机器人所在桌面建模，防止关节角碰到桌子
        T = aubo_FK(joint_point)
        link6_point = Point([T[5][0, 3], T[5][1, 3], T[5][2, 3]])
        link5_point = Point([T[4][0, 3], T[4][1, 3], T[4][2, 3]])
        link4_point = Point([T[3][0, 3], T[3][1, 3], T[3][2, 3]])
        link3_point = Point([T[2][0, 3], T[2][1, 3], T[2][2, 3]])
        link2_point = Point([T[1][0, 3], T[1][1, 3], T[1][2, 3]])
        if link2_point.z < self.limit_position or link3_point.z < self.limit_position or link4_point.z < self.limit_position or link5_point.z < self.limit_position or link6_point.z < self.limit_position:
            exceed_limit = True
            # print("关节角不符合要求")

        return exceed_limit

    def path_planning(self) -> List[Joint_point]:
        # print("length:",len(self.joint_point_list)) # length: 1
        count = 0  # Record the number of iterations

        start = time.time()
        while count < self.num_iterations:
            # Generate a new_radom_joint_point,生成一个0到1之间的随机数:random.random(), Xrand
            if random.random() > self.set_goal:
                new_random_joint_point = self.generate_new_joint_point()
            else:
                new_random_joint_point = self.goal

            # Get the nearest joint_point in self.joint_point_list,
            # 遍历所有节点,获得最近的距离所对应的索引
            index = self.get_the_nearest_joint_point(new_random_joint_point)
            # print("index",index)    # int型

            # Expand the tree
            nearest_joint_point = self.joint_point_list[index]
            new_joint_point = self.expand_new_joint_point(nearest_joint_point, new_random_joint_point)

            # Collision check，不需要写代码，我们只需要RRT生成的一个点即可
            if self.collision_check(new_joint_point) or self.check_joint_point_exceed_limit(new_joint_point):
                continue

            # Add the new joint_point into self.joint_point_list
            new_joint_point.signal = index  # 暂时不知道有啥用
            self.joint_point_list.append(new_joint_point)
            count = count + 1

        path = [self.get_farest_joint_point_from_start()]
        index = path[-1].signal
        while self.joint_point_list[index].signal is not None:
            joint_point = self.joint_point_list[index]
            path.append(joint_point)
            index = joint_point.signal
        path.append(self.start)
        path.reverse()
        # print("rrt path[-1]:",path[-1].theta)
        end = time.time()
        print("一次RRT耗费{:.2f}s".format(end - start))
        return path
